import cv2
import os
import numpy as np
from keras.optimizers import SGD
from keras.layers import Dense, Flatten, Dropout, Conv2D, MaxPooling2D
from keras.utils import np_utils
from sklearn.model_selection import train_test_split
from keras import Sequential
from keras import backend as K


class CaptureFace(object):
	def __init__(self):
		# 加载 级联分类器 模型：
		self.face_Cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
		self.eye_cascade = cv2.CascadeClassifier('haarcascade_eye.xml')
		self.faces = None
		self.image = None

	'''
	设备抓取照片
	'''
	def get_face(self):
		#  打开摄像头
		cap = cv2.VideoCapture(0)
		# 获取照片
		ret, self.image = cap.read()
		#  图像灰度化(降低运算强度)
		gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
		# 人脸检测 : faces中每一行中的元素分别表示检出的人脸在图中的（坐标x、坐标y、宽度、高度）
		# minSize是每个窗口的大小，minNeighbors会定义其周围有多少物体。
		self.faces = self.face_Cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5, minSize=(50, 50))
		# cv2.waitKey(0)
		cap.release()
		cv2.destroyAllWindows()

	'''
		给图片设置黑边框：
		max(宽,高),较短的一侧涂黑，使它变成和目标图像相同的比例，然后再resize，这样既可以保留原图的人脸信息，又可以防止图像形变
	'''
	def reset_size(self, image):
		size = (100, 100)
		_height, _width, _ = image.shape
		_edge = max(_height, _width)
		top, bottom, left, right = 0, 0, 0, 0
		if _height < _edge:
			_diff = _edge - _height
			top = int(_diff / 2)
			bottom = int(_diff - top)
		elif _width < _edge:
			_diff = _edge - _width
			left = int(_diff / 2)
			right = int(_diff - left)
		image_with_border = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])
		_image = cv2.resize(image_with_border, size)
		return _image

	def read_image(self):
		_x, _y = [], []
		for i in range(1, 177):
			try:
				im = cv2.imread('_imgs/%s.bmp' % str(i))
				if im is not None:
					_x.append(np.asarray(im, dtype=np.int8))
					_y.append(str(int((i-1)/11.0)))
			except Exception as es:
				print('error: ', es)
		return _x, _y


if __name__ == '__main__':
	K.set_image_dim_ordering('tf')
	##############################################################################################################
								# 读取所有照片并重新设置图片大小---------图片处理
	##############################################################################################################
	f = CaptureFace()
	RAW_IMAGE_DIR = 'imgs/'
	DATASET_DIR = '_imgs/'
	image_list = os.listdir(RAW_IMAGE_DIR)
	count = 1
	for _path in image_list:
		try:
			img = cv2.imread(os.path.join(RAW_IMAGE_DIR, _path))
			gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
			face = f.face_Cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5, flags=cv2.CASCADE_SCALE_IMAGE, minSize=(50, 50))
			if len(face) > 0:
				for (x, y, w, h) in face:
					roiImg = img[y:y + h, x:x + w]
					_img = f.reset_size(roiImg)
					cv2.imwrite('%s%d.bmp' % (DATASET_DIR, count), _img)
			count += 1
		except Exception as e:
			print('ERROR: ', e)

	##############################################################################################################
												# 训练数据划分
	##############################################################################################################
	'''
	卷积: Conv2D
	池化层: MaxPooling2D
	Dropout: 用来在训练时按一定概率随机丢弃一些神经元，以获得更高的训练速度以及防止过拟合
	Flatten: 用于卷积层与全连接层之间，把卷积输出的多维数据拍扁成一维数据送进全连接层（类似shape方法）
	梯度下降优化器: SGD
	'''
	# 读入所有图像及标签:
	raw_images, raw_labels = f.read_image()

	raw_images = np.asarray(raw_images, dtype=np.float64)
	raw_labels = np.asarray(raw_labels, dtype=np.int8)
	# 由于各标签的类别没有关系，使用one-hot编码规则，做到所有标签的平等化
	no_hot_label = np_utils.to_categorical(raw_labels)
	# 将所有数据划分成 训练集  测试集，训练集 ：测试集 = 7 : 3  ：
	train_input, valid_input, train_out, valid_out = train_test_split(raw_images, no_hot_label, test_size=0.3)

	# 数据归一化
	train_input /= 255.0
	valid_input /= 255.0
	##############################################################################################################
													# 训练
	##############################################################################################################
	# 构建卷积神经网络的每一层：
	'''
	卷积层参数:
		32个卷积核，每个卷积核是3 * 3，
		padding: 边缘不补充
		strides: 卷积步长向右、向下都为1,
		input_shape: 图片输入尺寸是（100，100， 3），
		激活函数: relu，也可以用sigmoid函数等，relu收敛速度比较快：
	'''
	# 选择训练模式
	_model = Sequential()
	# Keras提供了两套后端，th----->Theano:input_shape =(3,100, 100),tf----->tensorflow:input_shape =(100,10,3)
	input_shape = (100, 100, 3)
	_model.add(Conv2D(32, kernel_size=(3, 3), padding="valid", strides=(1, 1), input_shape=input_shape, activation="relu"))
	_model.add(Conv2D(32, kernel_size=(3, 3), padding="valid", strides=(1, 1), activation="relu"))

	# 池化层 过滤器尺寸是2 * 2
	_model.add(MaxPooling2D(pool_size=(2, 2)))
	# Dropout层
	_model.add(Dropout(0.2))

	_model.add(Conv2D(64, kernel_size=(3, 3), padding='valid', strides=(1, 1), activation='relu'))
	_model.add(Conv2D(64, kernel_size=(3, 3), padding='valid', strides=(1, 1), activation='relu'))
	_model.add(MaxPooling2D(pool_size=(2, 2)))
	_model.add(Dropout(0.2))
	# Flatten层，处于卷积层与Dense（全连层）之间，将图片的卷积输出压扁成一个一维向量：
	_model.add(Flatten())
	# 全连接层,  经典的神经网络结构，512个神经元：
	_model.add(Dense(512, activation='relu'))
	_model.add(Dropout(0.4))
	# 输出层，神经元数是标签种类数，使用sigmoid激活函数，输出最终结果：
	_model.add(Dense(len(no_hot_label[0]), activation='sigmoid'))
	# 神经网络结构打印
	# _model.summary()
	##############################################################################################################
										# SGD作为反向传播的优化器
	##############################################################################################################
	'''
	使用SGD作为反向传播的优化器，来使损失函数最小化，
	学习率(learning_rate)是0.01，
	学习率衰减因子(decay)用来随着迭代次数不断减小学习率，防止出现震荡。
	引入冲量(momentum),不仅可以在学习率较小的时候加速学习，又可以在学习率较大的时候减速，使用nesterov：
	'''
	learning_rate = 0.01
	decay = 1e-6
	momentum = 0.8
	nesterov = True
	_sgd = SGD(lr=learning_rate, decay=decay, momentum=momentum, nesterov=nesterov)
	# 编译模型，损失函数使用交叉熵，交叉熵函数随着输出和期望的差距越来越大，输出曲线会越来越陡峭，对权值的惩罚力度也会增大
	_model.compile(loss='categorical_crossentropy', optimizer=_sgd, metrics=['accuracy'])
	##############################################################################################################
											# 开始训练
	##############################################################################################################
	'''
	epochs: 训练100次（epochs），
	batch_size: 每次训练分几个批次，每批（batch_size）20个，
	shuffle: 用来打乱样本顺序：
	'''
	batch_size = 20
	epochs = 100
	_model.fit(train_input, train_out, epochs=epochs, batch_size=batch_size, shuffle=True, validation_data=(valid_input, valid_out))
	print(_model.evaluate(valid_input, valid_out, verbose=0))
	# 训练完成后在测试集上评估结果并保存模型供以后加载使用：
	MODEL_PATH = 'face_model.h5'
	_model.save(MODEL_PATH)
